feat(flexlb): batch scheduling infrastructure, strategy refactor, and error model hardening#1138
Open
wzy-99 wants to merge 41 commits into
Open
feat(flexlb): batch scheduling infrastructure, strategy refactor, and error model hardening#1138wzy-99 wants to merge 41 commits into
wzy-99 wants to merge 41 commits into
Conversation
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/1 · P1/0 · P2/1 · P3/0 Blocking IssuesP0
Non-blocking SuggestionsP2
Checklist Violations (3 fail / 56 total)General Principles Checklist
Python Static-First Checklist
Strengths
|
wzy-99
force-pushed
the
feat/flexlb_base
branch
2 times, most recently
from
5dff119 to
e5fa216
Compare
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/1 · P1/0 · P2/1 · P3/1 Blocking IssuesP0
Non-blocking SuggestionsP2
P3
Checklist Violations (7 fail / 69 total)General Principles Checklist
Python Static-First Checklist
Strengths
|
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/1 · P1/0 · P2/2 · P3/0 Blocking IssuesP0
Non-blocking SuggestionsP2
Checklist Violations (9 fail / 88 total)General Principles Checklist
RTP-LLM Checklist
Strengths
|
Collaborator
AI Code Review - PR #1138Status: LGTM Summary: P0/0 · P1/0 · P2/1 · P3/1 lgtm ready to ci Non-blocking SuggestionsP2
P3
Checklist Violations (7 fail / 103 total)General Principles Checklist
RTP-LLM Checklist
Strengths
|
|
CI dispatcher could not find a native This can happen if the PR was opened before the CI architecture change, or if the original run was deleted. To fix: push any commit (even empty: |
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/0 · P1/3 · P2/6 · P3/2 Blocking IssuesP1
Non-blocking SuggestionsP2
P3
Checklist Violations (1 fail / 56 total)General Principles Checklist
Strengths
|
…ng, and PrefillRpcServer metrics Combined from: - cc55836 feat: sync flexlb batch async fetch - b422c6c feat(flexlb): batch scheduling, strategy refactor, error model hardening, and test coverage - 89adee8 feat(PrefillRpcServer): enhance thread pool management with worker lambda pool and metrics - 909cfe5 feat(util): add requests library import for HTTP functionality
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/0 · P1/11 · P2/16 · P3/2 Blocking IssuesP1
Non-blocking SuggestionsP2
P3
Checklist Violations (7 fail / 56 total)General Principles Checklist
RTP-LLM Checklist
Strengths
|
…e + 3 P1 fixes - P0: move consecutiveFailures from GrpcWorkerStatusRunner instance field to WorkerStatus object, so the counter persists across sync cycles instead of being reset every time a new Runner is created - P1: fix isNonBatchPath to use != BATCH instead of == DIRECT, so AUTO-mode requests that fail batch qualification also get inflight tracking - P1: add null guard for newWorkerStatus.getRole().name() to prevent NPE crashing the sync thread - P1: add str branch in generate_config.py validate_role for JSON round-trip
Add per-request INFO-level event logs across engine/scheduler/executor/ dispatcher layers for complete request traceability (prefill → decode → finish). All events carry trace_id via streamLogTag() formatted as 'trace_id=XXX req_id=YYY'. Changes: - GenerateStream.h: implement streamLogTag() inline with trace_id+req_id - FIFOScheduler.cc: add request_activated (waiting→running + loading→ running) and request_finished (cleanup) logs - NormalExecutor.cc: add trace_id to prefill_batch_begin; add new decode_step_begin event for decode batch visibility - NormalOutputDispatcher.cc: add first_token + decode_finished per-request events (compatible with dispatchOutputAsync path) - frontend_server.py: add request_arrival + request_completion logs Design: all events use RTP_LLM_ACCESS_LOG_INFO (no [RANK][file:line] prefix); INFO only at batch boundaries/state transitions; trace_id coverage on all per-request events.
P1 #10: ResourceMeasure 重构导致 RandomStrategy 资源检查失效 - DecodeResourceMeasure/PrefillResourceMeasure 添加 @OverRide isResourceAvailable(WorkerEndpoint) 桥接方法 - 通过 instanceof 委托到类型化方法，修复之前仅走接口 default 方法（只查 isAlive()）的 bug P1 #8: activelyNotifyParticipants 通知条件反转 - 修正 localIp.equals 为 !localIp.equals，确保远程节点能收到 master 切换通知 P1 #9: waitForLeadershipTransfer 无退出上限 - 新增 MAX_WAIT_COUNT=30 最大等待保护，超时后 warn + 强制退出 P2 #3: applyTrafficPolicyOverride 缺 try-catch - JsonUtils.toObject 包裹 try-catch，防止异常 JSON 导致启动崩溃 P2 #12: LBStatusConsistencyService 静态线程池未关闭 - 添加 @PreDestroy shutdown() 关闭 SCHEDULED_EXECUTOR_SERVICE
…ia PDSepConfig - FlexlbConfig: flexlbBatchQueueMaxSize 64 → 1024 - PDSepConfig: add prefill_enqueue/worker_lambda/slot_pool_size fields (default 0 = formula) - PrefillRpcServer: initThreadPools() reads from pd_sep_config instead of hardcoded/env var - ConfigInit.cc: pybind registration + pickle for 3 new PDSepConfig fields - pd_separation_group_args.py: CLI args for 3 thread pool sizes
…n in .h) Commit f1df601 added an inline definition in GenerateStream.h but forgot to remove the existing out-of-line implementation in GenerateStream.cc, causing a redefinition error on cuda13_x86 builds.
- Add flexlbBatchFixedMaxInflightBatches config (default 0, disabled) to limit in-flight batches per prefill worker in fixed_window mode - Rename SLO batcher config flexlbBatchMaxInflightBatchesPerWorker -> flexlbBatchSloMaxInflightBatches for naming consistency - FixedWindowBatcherAlgorithm parks instead of dispatching when engine inflight batch count >= limit, preventing engine overload
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/0 · P1/11 · P2/44 · P3/0 Blocking IssuesP1
Non-blocking SuggestionsP2
Checklist ✅ (56 items passed)Strengths
|
…n multiple modules
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/0 · P1/12 · P2/22 · P3/3 Blocking IssuesP1
Non-blocking SuggestionsP2
P3
Checklist ✅ (56 items passed)Strengths
|
…ster routing Add MasterConfig.disable_domain_fallback field with env var MASTER_DISABLE_DOMAIN_FALLBACK to prevent fallback to VipServer domain routing when master is unavailable or not configured. When enabled, requests will fail with ROUTE_ERROR instead of silently degrading to domain-based service discovery.
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/1 · P1/15 · P2/31 · P3/2 Blocking IssuesP0
P1
Non-blocking SuggestionsP2
P3
Checklist Violations (6 fail / 56 total)General Principles Checklist
Strengths
|
… configurable default schedule mode
1. CostBasedPrefillStrategy.isNonBatchPath(): remove scheduleMode check
- When flexlbBatchEnabled=true, only check the global flag; schedule mode
is no longer a condition for creating request-ID placeholder entries.
- This eliminates ~93% redundant inflightBatches entries that were mixing
batch-ID and request-ID keys in the same ConcurrentHashMap.
2. PrefillEndpoint.calibrate(): log inflightBatches.size() instead of keySet()
- Avoid dumping 250+ IDs per calibrate log line when placeholder entries exist.
3. FlexlbConfig: add DEFAULT_SCHEDULE_MODE env var support
- New field defaultScheduleMode (String, default 'AUTO'), overridable via
DEFAULT_SCHEDULE_MODE env var (AUTO/BATCH/DIRECT).
4. FlexlbServiceImpl: use config-driven default schedule mode
- Replace hardcoded toScheduleMode() with resolveScheduleMode() that falls
back to config.getDefaultScheduleModeEnum() when proto sends default value.
…ng existing dashboard keys Add BatchSchedulerReporter that reports batcher queue depth, batch wait time, dispatch reason QPS, scheduler inflight size, and prefill inflight batch count. All 5 metric keys exactly match existing Grafana dashboard keys so batch-path data appears on the same panels as non-batch path without dashboard changes. Key changes: - BatchSchedulerReporter: 5 metrics (GAUGE/QPS, PRECISE) reusing ROUTING_QUEUE_*, ENGINE_BALANCING_MASTER_SELECT_DETAIL, ENGINE_LOCAL_TASK_MAP_SIZE, ENGINE_RUNNING_TASK_INFO_SIZE constants - MetricConstant: removed unused BATCH_* constants (now using existing keys) - FlexlbBatchScheduler.reportBatchMetrics: scheduler inflight + delegates per-worker metrics to PrefillEndpoint.reportBatchMetrics(reporter) - DefaultBatchDispatcher: cleanup removed gRPC timing metrics, empty reportGrpcFailure method, unused @scheduled import; reverted ThreadPoolExecutor back to ExecutorService - PrefillEndpoint: added reportBatchMetrics method (WorkerEndpoint state normalization principle) - Batcher components: wired BatchSchedulerReporter through constructor chain - Tests: updated constructor calls, 46 tests pass
…build sorted view The PriorityBlockingQueue iterator returns heap physical order (level-order traversal), not sorted order — so sortedItems() already needs copy+sort. The cache added a stale-view bug where WorkerBatcher.offer() adds items without invalidating the cache, causing processQueue to see incomplete views and split one logical batch into two dispatches (S->M) within 1-5ms. Queue size is typically <50 items, copy+sort cost is negligible (<1us) vs prefill latency (~200-500ms). Dropping the cache fixes the bug and simplifies the code.
When finishTask() is called after dequeue() has already moved the request from running_streams_ to finished_streams_, it created a duplicate entry with default batch_id=-1, polluting FlexLB's GetWorkerStatus reports. The root cause: finishTask always wrote to finished_streams_ even when the request was no longer in running_streams_. In FlexLB batch mode, this happens when prefill completes at remoteLoadCacheEnd (dequeue with batch_id=X) but remoteGenerate subsequently fails, triggering finishTask in finish_lambda which creates a second entry with batch_id=-1. Fix: return early from finishTask when the request is not found in running_streams_, since it was already reported via dequeue(). The decode failure is still propagated to FlexLB frontend through markResponseEntryDone, which triggers retry with a new batch_id.
Upgrade 5 key decode-side log statements to INFO level to enable request-level lifecycle tracing in production: - receive request (LocalRpcServer) - load cache from prefill done (DecodeRpcServer) - start to local generate (DecodeRpcServer) - local generate done (LocalRpcServer + DecodeRpcServer) Previously all decode inference events were DEBUG-only, making it impossible to trace a request through the decode pipeline at INFO log level. With this change, each request leaves a clear start-to-end trail on decode, matching the observability of the prefill side.
…l(20ms) to 5000ms The previous default fallback of syncRequestTimeoutMs to syncEngineStatusInterval (20ms) caused ALL GetCacheStatus gRPC calls to fail with DEADLINE_EXCEEDED, resulting in hit_cache=0 for all requests because WorkerStatus.cacheStatus was never populated. Set explicit default of 5000ms (5s) for gRPC calls to engine-side KV cache status queries, which typically take 600-800ms on L20B ARM architecture.
- StreamCacheResource::reuseCache() no longer requires per-request reuse_cache in protobuf; a global switch based on REUSE_CACHE=1 - DefaultBatchDispatcher sets reuse_cache=true & enable_device_cache=true as defensive measure for dispatched batch requests
1. GenerateStream: 将 FINISHED 状态置位前移到 releaseResource() 之前， 修复 Path B 中 tryReleaseKVBlock 因状态时序问题导致 insertIntoCache 从未被调用、SharedBlockCache items 恒为 0 的 bug 2. DefaultBatchDispatcher: 移除 buildInput 中对 reuse_cache / enable_device_cache 的硬编码 true 覆盖，改为从前端 protobuf 继承 3. DecodeResourceMeasure: 日志中 getIp() → ipPort() 提高可读性
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/0 · P1/4 · P2/8 · P3/1 Blocking IssuesP1
Non-blocking SuggestionsP2
P3
Checklist Violations (6 fail / 56 total)General Principles Checklist
Strengths
|
…gurable formula engine - New PrefillTimeFormula (flexlb-sync/strategy/): recursive descent parser supporting + - * / ^ (right-assoc), sqrt/log/exp/abs/max/min/pow - Rewrite PrefillTimePredictor: constructor takes formula string, estimateMs/predictBatchMs fill variable bindings then evaluate - FlexlbConfig: 6 costAlpha* fields replaced by single costFormula - ConfigService: PREFILL_COEFFICIENTS_ENV → PREFILL_TIME_FORMULA_ENV - Remove dead code: TaskInfo.estimatePrefillTime, old PrefillTimeFormula - Update all 5 test files; 23 tests pass
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/4 · P1/8 · P2/7 · P3/0 Blocking IssuesP0
P1
Non-blocking SuggestionsP2
Checklist Violations (8 fail / 56 total)General Principles Checklist
Strengths
|
…ache AND logic stopStream() no longer races with Engine Loop's advance() for KV cache persistence. Instead of actively calling finish() or reportError(), stopStream() now passively waits for Engine Loop to naturally detect GenerateDone → finish_internal() → FINISHED → insertIntoCache(). Changes: - stopStream(): when GenerateDone && !hasError, silently wait for Engine Loop to set FINISHED; only reportError for genuine errors - add prefill_stop_stream_wait_timeout_ms (default 2000ms) with timeout-based fallback to prevent indefinite wait - StreamCacheResource::reuseCache(): restore AND logic — both global REUSE_CACHE=1 AND per-request reuse_cache must be true - wire new config through Python CLI args, ConfigModules, and PrefillGenerateContext constructor in both GenerateStreamCall and EnqueueGroup paths This eliminates the race where stopStream()'s reportError(CANCELLED) could set hasError=true before finish_internal() could set FINISHED, blocking insertIntoCache and causing KV cache items=0 in PD prefill.
This reverts commit 5e46499.
LoadBalancer.rollBack(String ipPort, long requestId) → LoadBalancer.rollBack(WorkerEndpoint ep, long requestId) - DefaultRouter 注入 EndpointRegistry，在 rollBackRoutingFailure 中统一 查 endpoint 后再传给策略，策略不再各自持有 registry 引用 - RandomStrategy 彻底删除从未使用的 endpointRegistry 字段 - CostBasedPrefillStrategy/CostBasedDecodeStrategy 用 instanceof 替代 registry 查表，减少依赖 - 同步更新 4 个测试文件的构造器和 verify 调用 Tests: 48 passed, 0 failures
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/1 · P1/8 · P2/5 · P3/0 Blocking IssuesP0
P1
Non-blocking SuggestionsP2
Checklist Violations (8 fail / 56 total)General Principles Checklist
Python Static-First Checklist
Strengths
|
jianglan89
force-pushed
the
feat/flexlb_base
branch
from
b756fae to
20ce5ca
Compare
…y score Pass cacheHit and seqLen into computeScore to estimate per-request prefill time via PrefillTimePredictor. This makes cache match results directly participate in worker ranking, instead of being used only for SLO filtering and metrics reporting. score = predictor.estimateMs(seqLen, cacheHit) + batcherWaitMs() + realWaitTimeMs()
…lock key in updateCacheKeys FlexLB: WorkerStatus.updateFromResponse() unconditionally overwrote cacheStatus with resp.getCacheStatus(), which is null in GetWorkerStatus responses (cache status is only populated by GetCacheStatus gRPC). This nullified the cacheStatus set by GrpcCacheStatusCheckRunner on every status sync round, causing calculateCacheHit() to always return 0 even when GlobalCacheIndex had a successful prefix match — hit_cache was 0. Fix: only update cacheStatus when resp.getCacheStatus() is non-null. C++: updateCacheKeys() dropped the partial block key after incremental update and hard-coded lastBlockAligned=true. For sequences shorter than one block this produced zero cache keys, causing insertIntoCache to skip device cache insertion entirely. Fix: re-compute the partial block key when remaining tokens exist, and set lastBlockAligned consistent with initCacheKeys behavior.
FlexLB theory hit ratio (whale-lb.app.cache.theory.hit.ratio) was always 0 because cache_key_block_size was not included in FlexlbScheduleRequestPB, causing RecentCacheKeyTraceReporter.theoryHitTokens() to short-circuit at cacheKeyBlockSize <= 0. - Add cache_key_block_size field (field 13) to FlexlbScheduleRequestPB proto - Pass cache_key_block_size from Python master_client.py when building the request - Set cacheKeyBlockSize in FlexlbServiceImpl.buildContext() from proto field - Add test verifying cacheKeyBlockSize propagation through buildContext
…Server thread pools - Add PrefillPoolMetrics MetricsGroup (active/queued/completed/rejected/fallback/thread_max/queue_max) - Wire all 5 call sites with ScopeExit RAII guards for accurate active/completed tracking - Add reportOnePoolToKmonitor() called every 10s from GC thread - Demote periodic PoolMetrics log from INFO to DEBUG to avoid production noise - Add thread_max/queue_max to PoolMetrics for utilization monitoring
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/1 · P1/8 · P2/5 · P3/0 Blocking IssuesP0
P1
Non-blocking SuggestionsP2
Checklist Violations (14 fail / 56 total)General Principles Checklist
Python Static-First Checklist
Strengths
|
Problem: FlexlbBatchScheduler.cancel() only calls cancelPrefill (gRPC
cancel) and rollbackOnce (releases DecodeEndpoint resources), but never
removes the request from PrefillEndpoint.inflightBatches. When
FLEXLB_BATCH_FIXED_MAX_INFLIGHT_BATCHES is enabled, this causes the
inflight count to grow monotonically on cancels, triggering false
backpressure and starving the engine.
Changes in FlexlbBatchScheduler.java:
1. Added "volatile long batchId = -1" field to InflightEntry
- Stores the batch ID assigned by flushItems() so cancel() can
locate and repack the correct batch entry.
- Volatile ensures cross-thread visibility without synchronization.
- -1 sentinel means batch not yet committed; repack is skipped.
2. Set entry.batchId in flushItems() after commitBatch()
- Must be AFTER commitBatch (so repackBatch finds the entry in
PrefillEndpoint.inflightBatches) and BEFORE dispatch (so cancel
during gRPC can also repack).
3. Added repackPrefillBatch(InflightEntry) helper method
- Calls PrefillEndpoint.repackBatch(batchId, {requestId}) which:
- Single-request batch → removes entire entry (batch empty).
- Multi-request batch → keeps survivors, removes only this request.
- Batch already gone (calibrate/releaseBatch ran first) → no-op.
- Idempotent via ConcurrentHashMap.computeIfPresent.
4. Called repackPrefillBatch in three cancel paths:
- cancel(): main client-cancel path (timeout or upstream abort).
- onSuccess() cancelAfterAck: server reports cancel before ack.
- InflightEvictor callback: TTL eviction (default 300s) cleanup.
Edge cases handled:
- batchId=-1 (not yet committed): repack skipped, calibrate cleans up.
- Double repack (cancel + TTL): repackBatch is idempotent (no-op if
batch already removed).
- Concurrent repack + releaseBatch: ConcurrentHashMap serializes.
jianglan89
force-pushed
the
feat/flexlb_base
branch
from
b9e9b22 to
f990446
Compare
…, and decode inflight - Split mixed `app.engine.health.check.running.task.info.size` metric into: - `app.flexlb.batch.inflight.count` (prefill batch count, scheduler view) - `app.flexlb.batch.inflight.request.count` (prefill request count, scheduler view) - Add decode-side inflight metrics: - `app.flexlb.decode.inflight.count` (decode inflight requests, scheduler view) - `app.flexlb.decode.total.load` (decode total load: confirmed + inflight) - Add DecodeEndpoint.reportBatchMetrics() and EndpointRegistry.getDecodeEndpoints() - FlexlbBatchScheduler.reportBatchMetrics() now iterates both prefill and decode endpoints - BatchSchedulerReporter registers 13 metrics (was 10)
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/0 · P1/4 · P2/2 · P3/0 Blocking IssuesP1
auto task = std::make_shared<std::function<void()>>(std::move(finish_lambda));
auto error = slot_worker_pool_->pushTask([task] { (*task)(); });
if (error != autil::ThreadPoolBase::ERROR_NONE) {
...
std::thread fallback_thread([task] { (*task)(); });
fallback_thread.detach();
}或者在 pushTask 之前先 copy 一份
Non-blocking SuggestionsP2
PR 描述里应至少写出：动机、设计要点、rollout 顺序、兼容性影响、性能预期变化，方便 code review 与后续复盘。
Checklist Violations (8 fail / 56 total)General Principles Checklist
Strengths
|
…etrics reporting frequency
- Replace hardcoded lowercase role strings ('prefill'/'decode') with
RoleType.PREFILL.name()/RoleType.DECODE.name() across BatchSchedulerReporter,
PrefillEndpoint, FixedWindowBatcherAlgorithm, and FlexlbBatchScheduler
- Aligns scheduler-side metrics with EngineHealthReporter's role naming
convention (uppercase enum name) for consistent Grafana filtering
- Increase reportBatchMetrics frequency from 20s to 2s for finer granularity
…ch_end for one-shot extraction - begin: keep only ctx_batch/gen_batch/total_tokens/max_seq (brief marker) - end: add full streams detail + gen_batch/max_seq/forward_us (complete info) - add trace_id to MtpExecutor streams format (was missing, now matches NormalExecutor) - enables single grep prefill_batch_end to get batch composition + timing
|
Codex seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Collaborator
AI Code Review - PR #1138Status: BLOCKING Summary: P0/0 · P1/2 · P2/5 · P3/0 Blocking IssuesP1
Non-blocking SuggestionsP2
Checklist Violations (6 fail / 56 total)General Principles Checklist
Strengths
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR introduces FlexLB batch scheduling infrastructure with SLO-based admission control, strategy refactoring, and error model hardening.
Key Changes
Batch Scheduling Infrastructure
Strategy Refactor
Error Model Hardening (C++ engine)
Review Fixes (Round 3)
Test Coverage
Previous Reviews